/*
 * @Description: 闭环检测算法
 * @Author: Ren Qian
 * @Date: 2020-02-04 18:53:06
 */
#include "lidar_localization/mapping/loop_closing/loop_closing.hpp"

#include <cmath>
#include <algorithm>
#include <pcl/common/transforms.h>
#include <pcl/io/pcd_io.h>
#include "glog/logging.h"

#include "lidar_localization/global_defination/global_defination.h"
#include "lidar_localization/models/registration/ndt_registration.hpp"
#include "lidar_localization/models/cloud_filter/voxel_filter.hpp"
#include "lidar_localization/models/cloud_filter/no_filter.hpp"
#include "lidar_localization/tools/print_info.hpp"
#include <fstream>
#include "lidar_localization/tools/file_manager.hpp"

namespace lidar_localization {

template<class Matrix>
void write_matrix(const char* filename, const Matrix& matrix){
    std::ofstream out(filename, std::ios::out | std::ios::binary | std::ios::trunc);
    typename Matrix::Index rows=matrix.rows(), cols=matrix.cols();
    out.write((char*) (&rows), sizeof(typename Matrix::Index));
    out.write((char*) (&cols), sizeof(typename Matrix::Index));
    out.write((char*) matrix.data(), rows*cols*sizeof(typename Matrix::Scalar) );
    out.close();
}

template<class datatype>
void write_vector(const char* filename, const std::vector<datatype>& vec){
    std::ofstream out(filename, std::ios::out | std::ios::binary | std::ios::trunc);
    int num = vec.size();
    out.write((char*) (&num), sizeof(int));
    out.write((char*) vec.data(), num*sizeof(datatype) );
    out.close();
}

LoopClosing::LoopClosing() {
    InitWithConfig();
}

bool LoopClosing::InitWithConfig() {
    std::string config_file_path = WORK_SPACE_PATH + "/config/mapping/loop_closing.yaml";
    YAML::Node config_node = YAML::LoadFile(config_file_path);

    std::cout << "-----------------闭环检测初始化-------------------" << std::endl;
    InitParam(config_node);
    InitDataPath(config_node);
    InitRegistration(registration_ptr_, config_node);
    InitFilter("map", map_filter_ptr_, config_node);
    InitFilter("scan", scan_filter_ptr_, config_node);
    InitDetectionMethod(config_node);

    return true;
}

bool LoopClosing::InitParam(const YAML::Node& config_node) {
    extend_frame_num_ = config_node["extend_frame_num"].as<int>();
    loop_step_ = config_node["loop_step"].as<int>();
    diff_num_ = config_node["diff_num"].as<int>();
    detect_area_ = config_node["detect_area"].as<float>();
    fitness_score_limit_ = config_node["fitness_score_limit"].as<float>();

    return true;
}

bool LoopClosing::InitDataPath(const YAML::Node& config_node) {
    std::string data_path = config_node["data_path"].as<std::string>();
    if (data_path == "./") {
        data_path = WORK_SPACE_PATH;
    }

    key_frames_path_ = data_path + "/slam_data/key_frames";

    return true;
}

bool LoopClosing::InitRegistration(std::shared_ptr<RegistrationInterface>& registration_ptr, const YAML::Node& config_node) {
    std::string registration_method = config_node["registration_method"].as<std::string>();
    std::cout << "闭环点云匹配方式为：" << registration_method << std::endl;

    if (registration_method == "NDT") {
        registration_ptr = std::make_shared<NDTRegistration>(config_node[registration_method]);
    } else {
        LOG(ERROR) << "没找到与 " << registration_method << " 相对应的点云匹配方式!";
        return false;
    }

    return true;
}

bool LoopClosing::InitFilter(std::string filter_user, std::shared_ptr<CloudFilterInterface>& filter_ptr, const YAML::Node& config_node) {
    std::string filter_mothod = config_node[filter_user + "_filter"].as<std::string>();
    std::cout << "闭环的" << filter_user << "选择的滤波方法为：" << filter_mothod << std::endl;

    if (filter_mothod == "voxel_filter") {
        filter_ptr = std::make_shared<VoxelFilter>(config_node[filter_mothod][filter_user]);
    } else if (filter_mothod == "no_filter") {
        filter_ptr =std::make_shared<NoFilter>();
    } else {
        LOG(ERROR) << "没有为 " << filter_user << " 找到与 " << filter_mothod << " 相对应的滤波方法!";
        return false;
    }

    return true;
}

bool LoopClosing::InitDetectionMethod(const YAML::Node& config_node)
{
    detection_method_ = config_node["detection_method"].as<std::string>();

    if (detection_method_ == "scan_context")
    {
        use_yaw_ = config_node["scan_context"]["use_yaw"].as<bool>();
        scmanager_.LIDAR_HEIGHT = config_node["scan_context"]["LIDAR_HEIGHT"].as<double>(); // lidar height : add this for simply directly using lidar scan in the lidar local coord (not robot base coord) / if you use robot-coord-transformed lidar scans, just set this as 0.

        scmanager_.PC_NUM_RING = config_node["scan_context"]["PC_NUM_RING"].as<int>(); // 20 in the original paper (IROS 18)
        scmanager_.PC_NUM_SECTOR = config_node["scan_context"]["PC_NUM_SECTOR"].as<int>(); // 60 in the original paper (IROS 18)
        scmanager_.PC_MAX_RADIUS = config_node["scan_context"]["PC_MAX_RADIUS"].as<double>(); // 80 meter max in the original paper (IROS 18)
        scmanager_.PC_UNIT_SECTORANGLE = 360.0 / double(scmanager_.PC_NUM_SECTOR);
        scmanager_.PC_UNIT_RINGGAP = scmanager_.PC_MAX_RADIUS / double(scmanager_.PC_NUM_RING);

        // tree
        scmanager_.NUM_EXCLUDE_RECENT = config_node["scan_context"]["NUM_EXCLUDE_RECENT"].as<int>(); // simply just keyframe gap, but node position distance-based exclusion is ok. 
        scmanager_.NUM_CANDIDATES_FROM_TREE = config_node["scan_context"]["NUM_CANDIDATES_FROM_TREE"].as<int>(); // 10 is enough. (refer the IROS 18 paper)

        // loop thres
        scmanager_.SEARCH_RATIO = config_node["scan_context"]["SEARCH_RATIO"].as<double>(); // for fast comparison, no Brute-force, but search 10 % is okay. // not was in the original conf paper, but improved ver.
        // const double SC_DIST_THRES = 0.13; // empirically 0.1-0.2 is fine (rare false-alarms) for 20x60 polar context (but for 0.15 <, DCS or ICP fit score check (e.g., in LeGO-LOAM) should be required for robustness)
        scmanager_.SC_DIST_THRES = config_node["scan_context"]["SC_DIST_THRES"].as<double>(); // 0.4-0.6 is good choice for using with robust kernel (e.g., Cauchy, DCS) + icp fitness threshold / if not, recommend 0.1-0.15

        // config 
        scmanager_.TREE_MAKING_PERIOD_ = config_node["scan_context"]["TREE_MAKING_PERIOD"].as<int>(); // i.e., remaking tree frequency, to avoid non-mandatory every remaking, to save time cost / in the LeGO-LOAM integration, it is synchronized with the loop detection callback (which is 1Hz) so it means the tree is updated evrey 10 sec. But you can use the smaller value because it is enough fast ~ 5-50ms wrt N.

    }
    return true;
}

bool LoopClosing::Update(const KeyFrame key_frame, const KeyFrame key_gnss) {
    has_new_loop_pose_ = false;

    all_key_frames_.push_back(key_frame);
    if (detection_method_ == "gnss")
    {
        all_key_gnss_.push_back(key_gnss);
    }
    else if (detection_method_ == "scan_context")
    {
        // 生成当前scan
        CloudData::CLOUD_PTR scan_cloud_ptr(new CloudData::CLOUD());
        Eigen::Matrix4f scan_pose = Eigen::Matrix4f::Identity();
        JointScan(scan_cloud_ptr, scan_pose);
        scmanager_.makeAndSaveScancontextAndKeys(*scan_cloud_ptr);
        SaveScanContext();
    }
    else
    {
        return false;
    }
    

    int key_frame_index = 0;
    if (!DetectNearestKeyFrame(key_frame_index))
        return false;

    if (!CloudRegistration(key_frame_index))
        return false;

    has_new_loop_pose_ = true;
    return true;
}

bool LoopClosing::DetectNearestKeyFrame(int& key_frame_index) {
    static int skip_cnt = 0;
    static int skip_num = loop_step_;
    if (++skip_cnt < skip_num)
        return false;

    if ((int)all_key_frames_.size() < diff_num_ + 1)
        return false;

    if (detection_method_ == "gnss")
    {
        int key_num = (int)all_key_gnss_.size();
        float min_distance = 1000000.0;
        float distance = 0.0;

        KeyFrame history_key_frame;
        KeyFrame current_key_frame = all_key_gnss_.back();

        key_frame_index = -1;
        for (int i = 0; i < key_num - 1; ++i) {
            if (key_num - i < diff_num_)
                break;
            
            history_key_frame = all_key_gnss_.at(i);
            distance = fabs(current_key_frame.pose(0,3) - history_key_frame.pose(0,3)) + 
                    fabs(current_key_frame.pose(1,3) - history_key_frame.pose(1,3)) + 
                    fabs(current_key_frame.pose(2,3) - history_key_frame.pose(2,3));
            if (distance < min_distance) {
                min_distance = distance;
                key_frame_index = i;
            }
        }
        if (key_frame_index < extend_frame_num_)
            return false;

        skip_cnt = 0;
        skip_num = (int)min_distance;
        if (min_distance > detect_area_) {
            skip_num = std::max((int)(min_distance / 2.0), loop_step_);
            return false;
        } else {
            skip_num = loop_step_;
            return true;
        }
    }
    else if (detection_method_ == "scan_context")
    {
        auto res = scmanager_.detectLoopClosureID();
        key_frame_index = res.first;
        yaw_ = res.second;
        if (key_frame_index < extend_frame_num_)
            return false;
        
        skip_cnt = 0; 
        return true;
        
    }
    else
    {
        return false;
    }
    

}

bool LoopClosing::CloudRegistration(int key_frame_index) {
    // 生成地图
    CloudData::CLOUD_PTR map_cloud_ptr(new CloudData::CLOUD());
    Eigen::Matrix4f map_pose = Eigen::Matrix4f::Identity();
    JointMap(key_frame_index, map_cloud_ptr, map_pose);

    // 生成当前scan
    CloudData::CLOUD_PTR scan_cloud_ptr(new CloudData::CLOUD());
    Eigen::Matrix4f scan_pose = Eigen::Matrix4f::Identity();
    JointScan(scan_cloud_ptr, scan_pose);

    if (detection_method_ == "scan_context" && !use_yaw_)
    {
        scan_pose = all_key_frames_.at(key_frame_index).pose.inverse() * 
        all_key_frames_.back().pose;
        scan_pose.block<3, 1>(0, 3) = Eigen::Vector3f::Zero();
    }

    // 匹配
    Eigen::Matrix4f result_pose = Eigen::Matrix4f::Identity();
    Registration(map_cloud_ptr, scan_cloud_ptr, scan_pose, result_pose);

    // 计算相对位姿
    current_loop_pose_.pose = map_pose.inverse() * result_pose;

    // 判断是否有效
    if (registration_ptr_->GetFitnessScore() > fitness_score_limit_)
        return false;
    
    static int loop_close_cnt = 0;
    loop_close_cnt ++;

    std::cout << "检测到闭环 "<<  loop_close_cnt
              << ": 帧" << current_loop_pose_.index0 
              << "------>" << "帧" << current_loop_pose_.index1 << std::endl
              << "fitness score: " << registration_ptr_->GetFitnessScore() 
              << std::endl << std::endl;

    // std::cout << "相对位姿 x y z roll pitch yaw:";
    // PrintInfo::PrintPose("", current_loop_pose_.pose);

    return true;
}

bool LoopClosing::JointMap(int key_frame_index, CloudData::CLOUD_PTR& map_cloud_ptr, Eigen::Matrix4f& map_pose) {
    if (detection_method_ == "gnss")
    {
        map_pose = all_key_gnss_.at(key_frame_index).pose;
    }
    else if (detection_method_ == "scan_context")
    {
        // map_pose = all_key_frames_.at(key_frame_index).pose;
        map_pose = Eigen::Matrix4f::Identity();
    }
    else
    {
        return false;
    }
    
    
    current_loop_pose_.index0 = all_key_frames_.at(key_frame_index).index;
    
    // 合成地图
    Eigen::Matrix4f pose_to_gnss = map_pose * all_key_frames_.at(key_frame_index).pose.inverse();
    
    for (int i = key_frame_index - extend_frame_num_; i < key_frame_index + extend_frame_num_; ++i) {
        std::string file_path = key_frames_path_ + "/key_frame_" + std::to_string(all_key_frames_.at(i).index) + ".pcd";
        
        CloudData::CLOUD_PTR cloud_ptr(new CloudData::CLOUD());
        pcl::io::loadPCDFile(file_path, *cloud_ptr);
        
        Eigen::Matrix4f cloud_pose = pose_to_gnss * all_key_frames_.at(i).pose;
        pcl::transformPointCloud(*cloud_ptr, *cloud_ptr, cloud_pose);

        *map_cloud_ptr += *cloud_ptr;
    }
    map_filter_ptr_->Filter(map_cloud_ptr, map_cloud_ptr);
    return true;
}

bool LoopClosing::JointScan(CloudData::CLOUD_PTR& scan_cloud_ptr, Eigen::Matrix4f& scan_pose) {
    if (detection_method_ == "gnss")
    {
        scan_pose = all_key_gnss_.back().pose;
    }
    else if (detection_method_ == "scan_context")
    {
        // scan_pose = all_key_frames_.back().pose;
        scan_pose = Eigen::Matrix4f::Identity();
        if (use_yaw_)
        {
            Eigen::Affine3f icpInitialMatFoo = pcl::getTransformation(0, 0, 0, yaw_, 0, 0);
            scan_pose = icpInitialMatFoo.matrix();
        }
    }
    else
    {
        return false;
    }
    
    current_loop_pose_.index1 = all_key_frames_.back().index;
    current_loop_pose_.time = all_key_frames_.back().time;

    std::string file_path = key_frames_path_ + "/key_frame_" + std::to_string(all_key_frames_.back().index) + ".pcd";
    pcl::io::loadPCDFile(file_path, *scan_cloud_ptr);
    scan_filter_ptr_->Filter(scan_cloud_ptr, scan_cloud_ptr);

    return true;
}

bool LoopClosing::Registration(CloudData::CLOUD_PTR& map_cloud_ptr, 
                               CloudData::CLOUD_PTR& scan_cloud_ptr, 
                               Eigen::Matrix4f& scan_pose, 
                               Eigen::Matrix4f& result_pose) {
    // 点云匹配
    CloudData::CLOUD_PTR result_cloud_ptr(new CloudData::CLOUD());
    registration_ptr_->SetInputTarget(map_cloud_ptr);
    registration_ptr_->ScanMatch(scan_cloud_ptr, scan_pose, result_cloud_ptr, result_pose);

    return true;
}

bool LoopClosing::HasNewLoopPose() {
    return has_new_loop_pose_;
}

LoopPose& LoopClosing::GetCurrentLoopPose() {
    return current_loop_pose_;
}

bool LoopClosing::SaveScanContext()
{
    std::string file_path = key_frames_path_ + "/scan_context_" + std::to_string(all_key_frames_.back().index) + "_";
    write_matrix((file_path + "polarcontexts").c_str(), scmanager_.polarcontexts_.back());
    write_matrix((file_path + "polarcontext_invkeys").c_str(), scmanager_.polarcontext_invkeys_.back());
    write_matrix((file_path + "polarcontext_vkeys").c_str(), scmanager_.polarcontext_vkeys_.back());
    write_vector((file_path + "polarcontext_invkeys_mat").c_str(), scmanager_.polarcontext_invkeys_mat_.back());
    return true;
}

bool LoopClosing::SaveKeyFrames()
{
    std::ofstream odom_key_frames_ofs;
    if (!FileManager::CreateFile(odom_key_frames_ofs, key_frames_path_ + "/odom_key_frames.txt"))
        return false;

    for (size_t i = 0; i < all_key_frames_.size(); ++i) {
        SavePose(odom_key_frames_ofs, all_key_frames_[i].pose);
    }

    std::ofstream gnss_key_frames_ofs;
    if (!FileManager::CreateFile(gnss_key_frames_ofs, key_frames_path_ + "/gnss_key_frames.txt"))
        return false;

    for (size_t i = 0; i < all_key_gnss_.size(); ++i) {
        SavePose(gnss_key_frames_ofs, all_key_gnss_[i].pose);
    }

    return true;
}

bool LoopClosing::SavePose(std::ofstream& ofs, const Eigen::Matrix4f& pose) {
    for (int i = 0; i < 3; ++i) {
        for (int j = 0; j < 4; ++j) {
            ofs << pose(i, j);
            
            if (i == 2 && j == 3) {
                ofs << std::endl;
            } else {
                ofs << " ";
            }
        }
    }

    return true;
}

}